﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Xml;

namespace XmlDSigEx
{
    public class XmlC14NCanonicalisation : ICanonicalisationAlgorithm
    {
        internal const string URI_C14N = "http://www.w3.org/TR/2001/REC-xml-c14n-20010315";
        #region ICanonicalisationAlgorithm Members

        public Stream ProcessOctetStream(Stream stream, bool withComments)
        {
            // Must use XmlReader.Create to create a 'conformant' text reader.
            // See: http://msdn.microsoft.com/en-us/library/8459257a.aspx

            XmlReaderSettings settings = new XmlReaderSettings();
            settings.ProhibitDtd = false;
            settings.XmlResolver = null;

            XmlReader reader = XmlReader.Create(stream, settings);
            XmlDocument document = new XmlDocument() { PreserveWhitespace = true };
            document.Load(reader);
            string xpath = null;
            if (withComments)
            {
                xpath = "(//. | //@* | //namespace::*)";
            }
            else
            {
                xpath = "(//. | //@* | //namespace::*)[not(self::comment())]";
            }
            XmlNodeList nodes = document.SelectNodes(xpath);
            return ProcessNodeSet(nodes, withComments);
        }

        public Stream ProcessNodeSet(XmlNodeList nodeList, bool withComments)
        {
            List<XmlNode> nodeSet = new List<XmlNode>(nodeList.OfType<XmlNode>());
            return ProcessNodeSet(nodeSet, withComments);
        }

        public Stream ProcessNodeSet(List<XmlNode> nodeSet, bool withComments)
        {
            MemoryStream outputStream = new MemoryStream();
            StreamWriter outputWriter = new StreamWriter(outputStream);

            _ns_rendered = new Dictionary<string, string>();

            while (nodeSet.Count > 0)
            {
                XmlNode node = nodeSet[0];
                ProcessNode(node, nodeSet, outputWriter);
            }
            outputWriter.Flush();
            outputStream.Position = 0;
            return outputStream;
        }

        private void ProcessNode(XmlNode node, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            switch (node.NodeType)
            {
                case XmlNodeType.Document:
                    ProcessDocumentNode(node as XmlDocument, nodeSet, outputWriter);
                    break;

                case XmlNodeType.Element:
                    ProcessElement(node as XmlElement, nodeSet, outputWriter);
                    break;

                case XmlNodeType.ProcessingInstruction:
                    ProcessProcessingInstruction(node as XmlProcessingInstruction, nodeSet, outputWriter);
                    break;

                case XmlNodeType.Text:
                    ProcessTextNode(node as XmlText, nodeSet, outputWriter);
                    break;

                case XmlNodeType.Attribute:
                    ProcessAttributeOrNamespace(node as XmlAttribute, nodeSet, outputWriter);
                    break;

                case XmlNodeType.Whitespace:
                    ProcessWhitespace(node as XmlWhitespace, nodeSet, outputWriter);
                    break;

                case XmlNodeType.CDATA:
                    ProcessCData(node as XmlCDataSection, nodeSet, outputWriter);
                    break;

                case XmlNodeType.SignificantWhitespace:
                    ProcessSignificantWhitespace(node as XmlSignificantWhitespace, nodeSet, outputWriter);
                    break;

                default:
                    break;
            }
        }

        private void ProcessSignificantWhitespace(XmlSignificantWhitespace xmlSignificantWhitespace, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            if (nodeSet.Contains(xmlSignificantWhitespace))
            {
                if (xmlSignificantWhitespace.ParentNode != xmlSignificantWhitespace.OwnerDocument)
                {
                    outputWriter.Write(xmlSignificantWhitespace.Value);
                }
            }
        }

        private void ProcessCData(XmlCDataSection xmlCDataSection, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            outputWriter.Write(EncodeText(xmlCDataSection.Value));
        }

        private void ProcessWhitespace(XmlWhitespace xmlWhitespace, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            if (nodeSet.Contains(xmlWhitespace))
            {
                if (xmlWhitespace.ParentNode != xmlWhitespace.OwnerDocument)
                {
                    outputWriter.Write(xmlWhitespace.Value);
                }
            }
        }

        private void ProcessAttributeOrNamespace(XmlAttribute attributeNode, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            // BUG? Is this correct? How better to handle attributes in the node set but not part of an element?
            if (nodeSet.Contains(attributeNode.OwnerElement))
            {
                if (attributeNode.Name.StartsWith("xmlns"))
                {
                    ProcessNamespaceNode(attributeNode, nodeSet, attributeNode.OwnerElement, outputWriter);
                }
                else
                {
                    ProcessAttribute(attributeNode, nodeSet, outputWriter);
                }
            }
            else
            {
                nodeSet.Remove(attributeNode);
            }
        }

        private void ProcessAttribute(XmlAttribute attributeNode, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            outputWriter.Write(" ");
            outputWriter.Write(attributeNode.Name);
            outputWriter.Write("=");
            outputWriter.Write("\"");
            outputWriter.Write(ModifyAttributeValue(attributeNode.Value));
            outputWriter.Write("\"");
        }


        protected void ProcessNamespaceNode(XmlAttribute attributeNode, List<XmlNode> nodeSet, XmlElement ownerElement, StreamWriter outputWriter)
        {
            if (!IsXmlNamespaceNode(attributeNode))
            {
                if (ShouldWriteNamespaceNode(attributeNode, nodeSet, ownerElement))
                {
                    outputWriter.Write(" ");
                    outputWriter.Write(attributeNode.Name);
                    outputWriter.Write("=");
                    outputWriter.Write("\"");
                    outputWriter.Write(ModifyAttributeValue(attributeNode.Value));
                    outputWriter.Write("\"");

                    if (attributeNode.LocalName == "xmlns")
                    {
                        _ns_rendered[""] = attributeNode.Value;
                    }
                    else
                    {
                        _ns_rendered[attributeNode.LocalName] = attributeNode.Value;
                    }

                }
            }
        }

        protected virtual bool ShouldWriteNamespaceNode(XmlAttribute attributeNode, List<XmlNode> nodeSet, XmlElement ownerElement)
        {
            return !NearestIncludedAncestorDefines(attributeNode, nodeSet, ownerElement);
        }

        private bool NearestIncludedAncestorDefines(XmlAttribute attributeNode, List<XmlNode> nodeSet, XmlElement ownerElement)
        {
            XmlElement closestAncestor = FindNearestIncludedAncestor(ownerElement, nodeSet);
            if (closestAncestor != null)
            {
                XmlNodeList namespaces = closestAncestor.SelectNodes("namespace::*");
                foreach (XmlNode nsNode in namespaces)
                {
                    if (nsNode.LocalName == attributeNode.LocalName && nsNode.Value == attributeNode.Value)
                        return true;
                }
            }
            return false;
        }

        protected XmlElement FindNearestIncludedAncestor(XmlElement xmlElement, List<XmlNode> nodeSet)
        {
            XmlElement element = xmlElement;
            while (element.ParentNode as XmlElement != null)
            {
                if (nodeSet.Contains(element.ParentNode))
                    return element.ParentNode as XmlElement;
                element = element.ParentNode as XmlElement;
            }
            return null;
        }

        private bool IsXmlNamespaceNode(XmlAttribute attributeNode)
        {
            return attributeNode.LocalName == "xml" && attributeNode.Value == "http://www.w3.org/XML/1998/namespace";
        }

        private void ProcessTextNode(XmlText xmlText, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            outputWriter.Write(EncodeText(xmlText.Value));
        }

        private string EncodeText(string input)
        {
            input = input.Replace("&", "&amp;");
            input = input.Replace("<", "&lt;");
            input = input.Replace(">", "&gt;");
            input = input.Replace("\u000D", "&#xD;");
            return input;
        }

        private void ProcessDocumentNode(XmlDocument documentNode, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            foreach (XmlNode childNode in documentNode.ChildNodes)
            {
                ProcessNode(childNode, nodeSet, outputWriter);
            }

            foreach (XmlNode childNode in documentNode.ChildNodes)
            {
                nodeSet.Remove(childNode);
            }
            nodeSet.Remove(documentNode);
        }

        private void ProcessProcessingInstruction(XmlProcessingInstruction xmlProcessingInstruction, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            XmlDocument parentDocument = xmlProcessingInstruction.OwnerDocument;
            if (xmlProcessingInstruction.ParentNode == parentDocument)
            {
                List<XmlNode> documentChildren = new List<XmlNode>(parentDocument.ChildNodes.OfType<XmlNode>());
                if (documentChildren.IndexOf(xmlProcessingInstruction) > documentChildren.IndexOf(parentDocument.DocumentElement))
                {
                    outputWriter.Write("\u000A");
                }
            }

            outputWriter.Write("<?");
            outputWriter.Write(xmlProcessingInstruction.Target);
            if (!string.IsNullOrEmpty(xmlProcessingInstruction.Value))
            {
                outputWriter.Write(" ");
                outputWriter.Write(xmlProcessingInstruction.Value);
            }
            outputWriter.Write("?>");

            if (xmlProcessingInstruction.ParentNode == parentDocument)
            {
                List<XmlNode> documentChildren = parentDocument.ChildNodes.ConvertToListOf<XmlNode>();
                if (documentChildren.IndexOf(xmlProcessingInstruction) < documentChildren.IndexOf(parentDocument.DocumentElement))
                {
                    outputWriter.Write("\u000A");
                }
            }


            nodeSet.Remove(xmlProcessingInstruction);
        }

        protected Stack<Dictionary<string, string>> _renderedByOutputAncestor = new Stack<Dictionary<string, string>>();

        protected Dictionary<string, string> _ns_rendered;

        private void ProcessElement(XmlElement xmlElement, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            if (nodeSet.Contains(xmlElement))
            {

                Dictionary<string, string> newStateStack = new Dictionary<string, string>(_ns_rendered);
                _renderedByOutputAncestor.Push(newStateStack);
                _ns_rendered = new Dictionary<string, string>(newStateStack);

                outputWriter.Write("<");
                outputWriter.Write(xmlElement.Name);
                ProcessNamespaceAxis(xmlElement, nodeSet, outputWriter);
                ProcessAttributeAxis(xmlElement, nodeSet, outputWriter);
                outputWriter.Write(">");

            }
            ProcessChildren(xmlElement, nodeSet, outputWriter);
            if (nodeSet.Contains(xmlElement))
            {
                _ns_rendered = new Dictionary<string, string>(_renderedByOutputAncestor.Pop());
                outputWriter.Write("</");
                outputWriter.Write(xmlElement.Name);
                outputWriter.Write(">");
            }

            foreach (XmlAttribute attribute in xmlElement.Attributes)
            {
                nodeSet.Remove(attribute);
            }
            foreach (XmlNode childNode in xmlElement.ChildNodes)
            {
                nodeSet.Remove(childNode);
            }
            nodeSet.Remove(xmlElement);
        }

        private void ProcessChildren(XmlElement xmlElement, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            if (nodeSet.Contains(xmlElement))
            {
                foreach (XmlNode childNode in xmlElement.ChildNodes)
                {
                    ProcessNode(childNode, nodeSet, outputWriter);
                }
            }
            else
            {
                foreach (XmlElement childElement in xmlElement.ChildNodes.OfType<XmlElement>())
                {
                    ProcessElement(childElement, nodeSet, outputWriter);
                }
            }
        }

        private void ProcessAttributeAxis(XmlElement xmlElement, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            XmlNodeList attributeAxis = xmlElement.SelectNodes("attribute::*");
            List<XmlAttribute> attributes = new List<XmlAttribute>(attributeAxis.OfType<XmlAttribute>());

            // See Section 2.4 of spec.
            if (xmlElement.ParentNode != xmlElement.OwnerDocument && !nodeSet.Contains(xmlElement.ParentNode))
            {
                //if (xmlElement.ParentNode != xmlElement.OwnerDocument)
                //{
                MergeAncestorXmlAttributes(xmlElement, attributes);
            }

            attributes.Sort(new AttributeOrderComparer());
            foreach (XmlAttribute attributeNode in attributes)
            {
                ProcessAttribute(attributeNode, nodeSet, outputWriter);
            }
        }

        protected virtual void MergeAncestorXmlAttributes(XmlElement ownerElement, List<XmlAttribute> attributes)
        {
            List<XmlAttribute> ancestorList = new List<XmlAttribute>();
            XmlElement element = ownerElement.ParentNode as XmlElement;
            while (element != null)
            {
                var ancestorAttributes = element.SelectNodes("attribute::*").ConvertToListOf<XmlAttribute>();
                foreach (XmlAttribute ancestorAttribute in ancestorAttributes.Where(a => a.NamespaceURI == "http://www.w3.org/XML/1998/namespace"))
                {
                    if (!ancestorList.Exists(a => a.NamespaceURI == ancestorAttribute.NamespaceURI
                        && a.LocalName == ancestorAttribute.LocalName))
                    {
                        ancestorList.Add(ancestorAttribute);
                    }
                }
                element = element.ParentNode as XmlElement;
            }

            ancestorList.RemoveAll(a => attributes.Exists(b => b.LocalName == a.LocalName && b.NamespaceURI == a.NamespaceURI));
            attributes.AddRange(ancestorList);
        }

        private string ModifyAttributeValue(string attributeValue)
        {
            string value = attributeValue.Replace("&", "&amp;");
            value = value.Replace("<", "&lt;");
            value = value.Replace("\"", "&quot;");
            value = value.Replace("\u0009", "&#x9;");
            value = value.Replace("\u000A", "&#xA;");
            value = value.Replace("\u000D", "&#xD;");
            return value;
        }

        protected virtual void ProcessNamespaceAxis(XmlElement xmlElement, List<XmlNode> nodeSet, StreamWriter outputWriter)
        {
            XmlNodeList namespaceAxis = xmlElement.SelectNodes("namespace::*");
            List<XmlAttribute> namespaces = namespaceAxis.ConvertToListOf<XmlAttribute>();
            namespaces.RemoveAll(a => !nodeSet.Contains(a));
            namespaces.Sort(new NamespaceOrderComparer());

            // This block is to enforce the rules under section 2.3 of the spec concerning namespace axis.            
            if (namespaces.Count > 0 && namespaces[0].Name != "xmlns")
            {
                XmlElement nearestAncestor = FindNearestIncludedAncestor(xmlElement, nodeSet);
                if (nearestAncestor != null)
                {
                    XmlNodeList ancestorNamespaceAxis = nearestAncestor.SelectNodes("namespace::*");
                    // BUG! According to the spec, We should be checking that the ancestors node attribute
                    // should also be in the node set. However, it may have already been deleted from the nodeset.
                    // Need to fix.
                    if (ancestorNamespaceAxis
                        .ConvertToListOf<XmlAttribute>()
                        .Exists(a => a.Name == "xmlns"))
                    {
                        outputWriter.Write(" xmlns=\"\"");
                    }
                }
            }

            foreach (XmlAttribute namespaceNode in namespaces)
            {
                ProcessNamespaceNode(namespaceNode, nodeSet, xmlElement, outputWriter);
            }
        }

        #endregion
    }
}
